
from io import BytesIO
import plotly.graph_objects as go


def get_line_charts_image(x, yv_list, yn_list, text, layout_dict: dict = None):
    """
    折线图
    :param x: x轴数据列表，这里是日期列表
    :param yv_list: y轴数据列表,二维列表
    :param yn_list: y轴数据列表对应名字列表，一维列表
    :param text: y值对应文案
    :param layout_dict: 布局参数
    :return: 折线图io对象
    """
    if not x or not yv_list:
        fig = go.Figure(data=go.Scatter(x=[0], y=[0]))
    else:
        fig = go.Figure()
        for index, yv in enumerate(yv_list):
            if not isinstance(yv, list):
                continue
            fig.add_trace(go.Scatter(
                x=x,
                y=yv,
                text=text[index],
                textposition="top center",
                mode='lines+markers',
                name=yn_list[index])
            )
    if layout_dict:
        fig.update_layout(**layout_dict)
    image_io = BytesIO()
    fig.write_image(image_io, format="jpeg", scale=2)
    return image_io


def get_bar_charts_image(x, yv_list, yn_list, color_list=None, layout_dict=None):
    """
    柱状图
    :param x: x轴数据列表，这里是日期列表
    :param yv_list: y轴数据列表,二维列表
    :param yn_list: y轴数据列表对应名字列表，一维列表
    :param color_list: y轴数据柱状图颜色
    :param layout_dict: 布局参数
    :return: 柱状图io对象
    """
    if not x or not yv_list:
        fig = go.Figure(data=go.Bar(x=[0], y=[0]))
    else:
        fig = go.Figure()
        for index, yv in enumerate(yv_list):
            if not isinstance(yv, list):
                continue
            item_name = yn_list[index]
            marker = None
            if color_list:
                marker = {'color': color_list[index]}
            fig.add_trace(
                go.Bar(x=x, y=yv, name=item_name, text=yv,
                       textfont={'size': 18, 'color': 'rgba(45,53,67,65)'},
                       textposition='outside', marker=marker)
            )
    legend = {
        'title_font_family': 'ABeeZee',
        'font': {'family': 'ABeeZee', 'size': 12, 'color': 'rgba(45,53,67,65)'},
        'traceorder': 'reversed',
        'y': 1.04,
        'x': 0,
        'bgcolor': 'rgba(0,0,0,0)',
        'orientation': 'h'
    }
    xaxis = {
        "tickfont": {"size": 16},
    }
    yaxis = {
        'gridcolor': 'rgba(65,97,128,15)',
        'tick0': 0,
        'zeroline': True,
        'zerolinewidth': 1,
        'zerolinecolor': 'rgba(65,97,128,15)'
    }
    if len(x) <= 2:
        fig.update_layout(bargap=0.8)
    elif len(x) <= 5:
        fig.update_layout(bargap=0.5)
    layout = dict(
        paper_bgcolor='rgba(255,255,255,100)',
        plot_bgcolor='rgba(255,255,255,100)',
        margin={'l': 5, 'r': 5, 't': 35, 'b': 5},
        xaxis=xaxis,
        yaxis=yaxis,
        legend=legend
    )
    if layout_dict:
        layout.update(layout_dict)
    fig.update_layout(**layout)
    image_io = BytesIO()
    fig.write_image(image_io, format="jpeg", scale=2)
    return image_io


def get_pie_charts_image(text, values, labels=None, title=None, hole=0.5, color_list=None):
    """
    生成饼图
    Args:
        text: 标签序列
        values: 数据序列
        title: 标题
        labels: 图例文案
        hole: 空心
        color_list: 颜色
    Returns:饼图IO对象

    """
    draw_dict = dict(
        values=values,
        hole=hole,
        text=text,
        textposition="inside",
        labels=labels
    )
    if color_list:
        draw_dict.update({"marker": {"colors": color_list}})
    fig = go.Figure(
        data=[go.Pie(**draw_dict)]
    )
    fig.update_layout(
        legend=dict(x=0.9, y=1),
        title=title,
        title_x=0.5,
        title_y=0.95,
    )
    image_io = BytesIO()
    fig.write_image(image_io, format="jpeg", scale=2)
    return image_io
